import type { Plugin } from 'vite';
import { parse } from '@babel/parser';
import traverse, { NodePath } from '@babel/traverse';
import generate from '@babel/generator';
import * as t from '@babel/types';

const realGenerate = (generate as any).default || generate;

export function devtools(): Plugin {
  return {
    name: 'vite-plugin-react-data-source',
    enforce: 'pre',
    async transform(code: string, id: string) {
      if (!id.endsWith('.tsx') && !id.endsWith('.jsx')) return;

      const ast = parse(code, {
        sourceType: 'module',
        plugins: ['jsx', 'typescript'],
      });

      // 兼容 ESM/CJS 的 traverse
      const realTraverse = (traverse as any).default || traverse;

      realTraverse(ast, {
        ReturnStatement(path: NodePath<t.ReturnStatement>) {
          const argument = path.node.argument;
          if (
            argument &&
            (t.isJSXElement(argument) || t.isJSXFragment(argument))
          ) {
            if (t.isJSXElement(argument)) {
              const hasDataSource = argument.openingElement.attributes.some(
                (attr) =>
                  t.isJSXAttribute(attr) &&
                  attr.name.name === 'data-source'
              );
              if (!hasDataSource) {
                // 组件名
                let funcPath = path.findParent((p) =>
                  p.isFunctionDeclaration() ||
                  p.isFunctionExpression() ||
                  p.isArrowFunctionExpression()
                ) as NodePath<any> | null;
                let name = 'Anonymous';
                if (funcPath && funcPath.node.id && funcPath.node.id.name) {
                  // 尝试从函数节点获取函数名
                  name = funcPath.node.id.name;
                } else if (funcPath) {
                  // 尝试从父节点获取变量名
                  const parent = funcPath.parentPath;
                  if (parent) {
                    if (parent.isVariableDeclarator() && parent.node.id && parent.node.id.type === 'Identifier') {
                      name = parent.node.id.name;
                    }
                    else if (parent.isAssignmentExpression() && parent.node.left && parent.node.left.type === 'Identifier') {
                      name = parent.node.left.name;
                    }
                  }
                }
                const root = process.cwd().replace(/\\/g, '/');
                let relFile = id.replace(/\\/g, '/');
                if (relFile.startsWith(root)) {
                  relFile = relFile.slice(root.length);
                  if (relFile.startsWith('/')) relFile = relFile.slice(1);
                }
                let source = '';
                if (funcPath && funcPath.node.start != null && funcPath.node.end != null) {
                  source = code.slice(funcPath.node.start, funcPath.node.end);
                }
                argument.openingElement.attributes.push(
                  t.jsxAttribute(
                    t.jsxIdentifier('data-source'),
                    t.jsxExpressionContainer(
                      t.stringLiteral(JSON.stringify({ name, file: relFile, source }))
                    )
                  )
                );
              }
            }
          }
        },
      });

      const output = realGenerate(ast, {}, code);
      return {
        code: output.code,
        map: null,
      };
    },
  };
}